#include "Graphics/PiccoloVKDescriptorSet.h"
#include "Graphics/PiccoloVKDevice.h"

namespace Piccolo
{
    /////////////////////////////////////////// PiccoloVKDescriptorSetLayout /////////////////////////////////////////////////////

    PiccoloVKDescriptorSetLayout::PiccoloVKDescriptorSetLayout(PiccoloVKDevice *device, const std::vector<VkDescriptorSetLayoutBinding> &bindings) : mDevice(device)
    {
        VkDescriptorSetLayoutCreateInfo descriptorSetLayoutInfo =
        {
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .bindingCount = static_cast<uint32_t>(bindings.size()),
            .pBindings = bindings.data()
        };
        CALL_VK(vkCreateDescriptorSetLayout(mDevice->GetLogicDevice(), &descriptorSetLayoutInfo, nullptr, &mDescriptorSetLayout));
    }

    PiccoloVKDescriptorSetLayout::~PiccoloVKDescriptorSetLayout()
    {
        VK_D(DescriptorSetLayout, mDevice->GetLogicDevice(), mDescriptorSetLayout);
    }

    /////////////////////////////////////////// PiccoloVKDescriptorSetLayout Over /////////////////////////////////////////////////////


    /////////////////////////////////////////// PiccoloVKDescriptorPool /////////////////////////////////////////////////////

    PiccoloVKDescriptorPool::PiccoloVKDescriptorPool(PiccoloVKDevice *device, uint32_t maxSets, const std::vector<VkDescriptorPoolSize> &poolSizes) : mDevice(device)
    {
        VkDescriptorPoolCreateInfo descriptorPoolInfo =
        {
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .maxSets = maxSets,
            .poolSizeCount = static_cast<uint32_t>(poolSizes.size()),
            .pPoolSizes = poolSizes.data()
        };
        CALL_VK(vkCreateDescriptorPool(mDevice->GetLogicDevice(), &descriptorPoolInfo, nullptr, &mDescriptorPool));
    }

    PiccoloVKDescriptorPool::~PiccoloVKDescriptorPool()
    {
        VK_D(DescriptorPool, mDevice->GetLogicDevice(), mDescriptorPool);
    }

    std::vector<VkDescriptorSet> PiccoloVKDescriptorPool::AllocateDescriptorSet(PiccoloVKDescriptorSetLayout* descriptorSetLayout, uint32_t count)
    {
        std::vector<VkDescriptorSet> descriptorSets(count);
        std::vector<VkDescriptorSetLayout> descriptorSetLayouts(count);
        for (int i = 0; i < count; i++)
        {
            descriptorSetLayouts[i] = descriptorSetLayout->GetDescriptorSetLayout();
        }

        VkDescriptorSetAllocateInfo allocateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
            .pNext = nullptr,
            .descriptorPool = mDescriptorPool,
            .descriptorSetCount = count,
            .pSetLayouts = descriptorSetLayouts.data()
        };
        VkResult ret = vkAllocateDescriptorSets(mDevice->GetLogicDevice(), &allocateInfo, descriptorSets.data());
        CALL_VK(ret);
        if (ret != VK_SUCCESS)
        {
            descriptorSets.clear();
        }
        return descriptorSets;
    }

    /////////////////////////////////////////// PiccoloVKDescriptorPool Over /////////////////////////////////////////////////////

    /////////////////////////////////////////// DescriptorSetWriter /////////////////////////////////////////////////////

    VkDescriptorBufferInfo DescriptorSetWriter::BuildBufferInfo(VkBuffer buffer, VkDeviceSize offset, VkDeviceSize range)
    {
        VkDescriptorBufferInfo bufferInfo
        {
            .buffer = buffer,
            .offset = offset,
            .range = range
        };
        return bufferInfo;
    }

    VkDescriptorImageInfo DescriptorSetWriter::BuildImageInfo(VkSampler sampler, VkImageView imageView)
    {
        VkDescriptorImageInfo imageInfo
        {
            .sampler = sampler,
            .imageView = imageView,
            .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
        };
        return imageInfo;
    }

    VkWriteDescriptorSet DescriptorSetWriter::WriteBuffer(VkDescriptorSet dstSet, uint32_t dstBinding, VkDescriptorType descriptorType, VkDescriptorBufferInfo *pBufferInfo)
    {
        VkWriteDescriptorSet writeDescriptorSet
        {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .pNext = nullptr,
            .dstSet = dstSet,
            .dstBinding = dstBinding,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = descriptorType,
            .pBufferInfo = pBufferInfo,
        };
        return writeDescriptorSet;
    }

    VkWriteDescriptorSet DescriptorSetWriter::WriteImage(VkDescriptorSet dstSet, uint32_t dstBinding, VkDescriptorType descriptorType, VkDescriptorImageInfo *pImageInfo)
    {
        VkWriteDescriptorSet writeDescriptorSet
        {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .pNext = nullptr,
            .dstSet = dstSet,
            .dstBinding = dstBinding,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = descriptorType,
            .pImageInfo = pImageInfo,
        };
        return writeDescriptorSet;
    }

    void DescriptorSetWriter::UpdateDescriptorSets(VkDevice device, const std::vector<VkWriteDescriptorSet> &writes)
    {
        vkUpdateDescriptorSets(device, writes.size(), writes.data(), 0, nullptr);
    }

    /////////////////////////////////////////// DescriptorSetWriter Over /////////////////////////////////////////////////////
}
